/******************************************************************************
	This file is a part of the Snapshot Bioinformatics Project

	Copyright (C) 2012-2013 Jarom Schow

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************/

#include "TaxonUpdateTree.h"

#include <QtCore/QTextStream>
#include <QtCore/QStack>

#include <QtSql/QSqlQuery>

#include "Core.h"
#include "Progress.h"
#include "SqlServer.h"
#include "SqlInsertBuilder.h"
#include "SqlUpdateBuilder.h"

using namespace Bio;

/**
**/
int	TaxonUpdateNode::s_taxaInserted = 0;

/**
**/
int TaxonUpdateNode::s_taxaUpdated = 0;

/**
**/
TaxonUpdateNode::TaxonUpdateNode() : m_taxId( -1 ), m_parentId( -1 ), m_genCodeId( -1 ), m_mitoGenCodeId( -1 ), m_localId( -1 ),
	m_localParentId( -1 ), m_firstChild( NULL ), m_lastChild( NULL ), m_nextSibling( NULL ), m_left( -1 ), m_right( -1 ), 
	m_flags( UpdateInsert )
{
}

/**
**/
TaxonUpdateNode::TaxonUpdateNode( int taxId, int parentId, const QString &rank, int genCodeId, int mitoGenCodeId ) : m_taxId( taxId ),
	m_parentId( parentId ), m_rank( rank ), m_genCodeId( genCodeId ), m_mitoGenCodeId( mitoGenCodeId ), m_localId( -1 ),
	m_localParentId( -1 ), m_firstChild( NULL ), m_lastChild( NULL ), m_nextSibling( NULL ), m_left( -1 ), m_right( -1 ), 
	m_flags( UpdateInsert )
{
}

/**
**/
TaxonUpdateNode::TaxonUpdateNode( int localId, int taxId, int localParentId, const QString &rank, int genCodeId,
						int mitoGenCodeId, int left, int right ) : 
	m_localId( localId ), m_taxId( taxId ), m_localParentId( localParentId ), m_rank( rank ), m_genCodeId( genCodeId ),
	m_mitoGenCodeId( mitoGenCodeId ), m_parentId( -1 ), m_firstChild( NULL ), m_lastChild( NULL ), m_nextSibling( NULL ),
	m_left( left ), m_right( right ), m_flags( UpdateNone )
{
}

/**
**/
TaxonUpdateNode::~TaxonUpdateNode()
{
	TaxonUpdateNode	*child = m_firstChild;

	while( child ){
		TaxonUpdateNode	*next = child->m_nextSibling;

		delete child;
		child = next;
	}

	m_firstChild = NULL;
	m_lastChild = NULL;
	m_nextSibling = NULL;
}

/**
**/
TaxonUpdateNode* TaxonUpdateNode::firstChild() const
{
	return m_firstChild;
}

/**
**/
TaxonUpdateNode* TaxonUpdateNode::nextSibling() const
{
	return m_nextSibling;
}

/**
**/
TaxonUpdateNode* TaxonUpdateNode::takeFirstChild()
{
	TaxonUpdateNode	*child = m_firstChild;

	if( m_firstChild ){
		m_firstChild = child->m_nextSibling;
		child->m_nextSibling = NULL;
		if( !m_firstChild )
			m_lastChild = NULL;
	}

	return child;
}

/**
**/
void TaxonUpdateNode::addChild( TaxonUpdateNode *child )
{
	if( !child )
		return;

	if( m_lastChild ){
		m_lastChild->m_nextSibling = child;
		m_lastChild = child;
	}
	else {
		m_firstChild = child;
		m_lastChild = child;
	}
}

/**
**/
int TaxonUpdateNode::getNumChildren( bool recursive )
{
	int				nChildren = 0;
	TaxonUpdateNode	*child = m_firstChild;

	while( child ){
		nChildren++;
		if( recursive )
			nChildren += child->getNumChildren( true );
		child = child->m_nextSibling;
	}

	return nChildren;
}

/**
**/
void TaxonUpdateNode::setTaxonId( int id )
{
	m_taxId = id;
}

/**
**/
int TaxonUpdateNode::getTaxonId() const
{
	return m_taxId;
}

/**
**/
void TaxonUpdateNode::setParentId( int id )
{
	m_parentId = id;
}

/**
**/
int TaxonUpdateNode::getParentId() const
{
	return m_parentId;
}

/**
**/
void TaxonUpdateNode::setGeneticCode( int id )
{
	m_genCodeId = id;
}

/**
**/
int TaxonUpdateNode::getGeneticCode() const
{
	return m_genCodeId;
}

/**
**/
void TaxonUpdateNode::setMitoGeneticCode( int id )
{
	m_mitoGenCodeId = id;
}

/**
**/
int TaxonUpdateNode::getMitoGeneticCode() const
{
	return m_mitoGenCodeId;
}

/**
**/
void TaxonUpdateNode::setLocalId( int id )
{
	m_localId = id;
}

/**
**/
int TaxonUpdateNode::getLocalId() const
{
	return m_localId;
}

/**
**/
void TaxonUpdateNode::setLocalParentId( int id )
{
	m_localParentId = id;
}

/**
**/
int TaxonUpdateNode::getLocalParentId() const
{
	return m_localParentId;
}

/**
**/
void TaxonUpdateNode::setRank( const QString &rank )
{
	m_rank = rank;
}

/**
**/
QString TaxonUpdateNode::getRank() const
{
	return m_rank;
}

/**
**/
void TaxonUpdateNode::update( const QString &rank, int genCodeId, int mitoGenCodeId )
{
	if( m_rank != rank ){
		m_rank = rank;
		m_flags |= UpdateRank;
	}

	if( m_genCodeId != genCodeId ){
		m_genCodeId = genCodeId;
		m_flags |= UpdateGenCode;
	}

	if( m_mitoGenCodeId != mitoGenCodeId ){
		m_mitoGenCodeId = mitoGenCodeId;
		m_flags |= UpdateMitoGenCode;
	}
}

/**
**/
void TaxonUpdateNode::finalize( int nextTaxonId )
{
	int	rightLeft = 1;

	finalizeRecurse( nextTaxonId, rightLeft );

	// the root parent id points to itself
	m_localParentId = m_localId;
}

/**
**/
bool TaxonUpdateNode::updateDB( SqlServer &server, Progress *progress )
{
	QString				sql( "INSERT INTO taxon (taxon_id, ncbi_taxon_id, "
							"parent_taxon_id, node_rank, genetic_code, mito_genetic_code, "
							"left_value, right_value) VALUES " );
	SqlInsertBuilder	insert( sql, &server );

	s_taxaInserted = 0;
	s_taxaUpdated = 0;

	if( !updateDBRecurse( server, insert, progress ) )
		return false;

	Bio::Core::log( "Taxon update finished." );
	Bio::Core::log( QString( "\tInserted: %1   Updated: %2" ).arg( s_taxaInserted ).arg( s_taxaUpdated ) );

	return insert.flush();
}

/**
**/
void TaxonUpdateNode::print( QTextStream &os )
{
	printRecurse( os, 0 );
}

/**
**/
void TaxonUpdateNode::finalizeRecurse( int &nextTaxonId, int &rightLeft )
{
	if( m_localId < 0 ){
		m_localId = nextTaxonId++;
		m_flags |= UpdateTaxonId;
	}

	if( m_left != rightLeft ){
		m_left = rightLeft;
		m_flags |= UpdateLeft;
	}

	rightLeft++;
	
	TaxonUpdateNode	*child = m_firstChild;

	while( child ){
		child->setLocalParentId( m_localId );
		child->finalizeRecurse( nextTaxonId, rightLeft );
		child = child->m_nextSibling;
	}

	if( m_right != rightLeft ){
		m_right = rightLeft;
		m_flags |= UpdateRight;
	}

	rightLeft++;
}

/**
**/
bool TaxonUpdateNode::updateDBRecurse( SqlServer &server, SqlInsertBuilder &insert, Progress *progress )
{
	if( m_flags != UpdateNone ){
		if( m_flags & UpdateInsert ){
			if( !insertInDB( insert ) )
				return false;
		}
		else {
			if( !updateDBInternal( server ) )
				return false;
		}
	}

	TaxonUpdateNode	*child = m_firstChild;

	while( child ){
		if( !child->updateDBRecurse( server, insert, progress ) )
			return false;

		child = child->m_nextSibling;
	}

	if( progress ){
		progress->step();
		if( progress->isCanceled() )
			return false;
	}

	return true;
}

/**
**/
bool TaxonUpdateNode::updateDBInternal( SqlServer &server )
{
	QString				condition = "WHERE taxon_id = " + QString::number( m_localId );
	SqlUpdateBuilder	update( "taxon", condition, &server );

	if( m_flags & UpdateParentId )
		update.addValue( "parent_taxon_id", m_parentId );

	if( m_flags & UpdateGenCode )
		update.addValue( "genetic_code", m_genCodeId );

	if( m_flags & UpdateMitoGenCode )
		update.addValue( "mito_genetic_code", m_mitoGenCodeId );

	if( m_flags & UpdateRank )
		update.addValue( "node_rank", m_rank );

	if( m_flags & UpdateLeft )
		update.addValue( "left_value", m_left );

	if( m_flags & UpdateRight )
		update.addValue( "right_value", m_right );
	
	return update.exec();
}

/**
**/
bool TaxonUpdateNode::insertInDB( SqlInsertBuilder &insert )
{
	insert.addValue( m_localId );
	insert.addValue( m_taxId );
	insert.addValue( m_localParentId );
	insert.addValue( m_rank );
	insert.addValue( m_genCodeId );
	insert.addValue( m_mitoGenCodeId );
	insert.addValue( m_left );
	insert.addValue( m_right );

	return insert.addRow();
}

/**
**/
void TaxonUpdateNode::printRecurse( QTextStream &os, int level )
{
	for( int i = 0; i < level; i++ )
		os << '\t';

	os << "TID: " << m_taxId;
	os << "\tPID: " << m_parentId;
	//os << "\tGC: " << m_genCodeId;
	//os << "\tMGC: " << m_mitoGenCodeId;
	os << "\tLID: " << m_localId;
	os << "\tRank: " << m_rank;
	os << "\tLeft: " << m_left;
	os << "\tRight: " << m_right;
	os << "\n";

	TaxonUpdateNode	*child = m_firstChild;

	while( child ){
		child->printRecurse( os, level + 1 );
		child = child->m_nextSibling;
	}
}

/**
**/
TaxonUpdateName::TaxonUpdateName()
{
}

/**
**/
void TaxonUpdateName::addName( const QString &name, const QString &nameClass )
{
	m_names.append( QPair<QString,QString>( name, nameClass ) );
}

/**
**/
int TaxonUpdateName::getNumNames() const
{
	return m_names.count();
}

/**
**/
QString TaxonUpdateName::getName( int which ) const
{
	if( which < 0 || which >= m_names.count() )
		return QString::null;

	return m_names[which].first;
}

/**
**/
QString TaxonUpdateName::getNameClass( int which ) const
{
	if( which < 0 || which >= m_names.count() )
		return QString::null;

	return m_names[which].second;
}

/**
**/
QListIterator< QPair<QString,QString> > TaxonUpdateName::getNameIterator() const
{
	return QListIterator< QPair<QString,QString> >( m_names );
}

/**
**/
TaxonUpdateTree::TaxonUpdateTree() : m_root( NULL ), m_orphans( NULL )
{
}

/**
**/
TaxonUpdateTree::~TaxonUpdateTree()
{
	delete m_root;
	m_root = NULL;

	delete m_orphans;
	m_orphans = NULL;
}

/**
**/
bool TaxonUpdateTree::insertNode( int taxId, int parentId, const QString &rank, int genCodeId, int mitoGenCodeId )
{
	if( m_lookupTable.contains( taxId ) ){
		m_lookupTable[taxId]->update( rank, genCodeId, mitoGenCodeId );
		return true;
	}

	TaxonUpdateNode	*node = new TaxonUpdateNode( taxId, parentId, rank, genCodeId, mitoGenCodeId );

	if( m_lookupTable.contains( parentId ) )
		m_lookupTable[parentId]->addChild( node );
	else if( !m_root )
		m_root = node;
	else {
		if( !m_orphans )
			m_orphans = new TaxonUpdateNode;
		m_orphans->addChild( node );
	}

	m_lookupTable[taxId] = node;

	return true;
}

/**
**/
bool TaxonUpdateTree::insertName( int taxId, const QString &name, const QString &nameClass )
{
	m_names[taxId].addName( name, nameClass );

	return true;
}

/**
**/
bool TaxonUpdateTree::finalize( SqlServer &server )
{
	if( m_orphans ){
		TaxonUpdateNode	*node = m_orphans->takeFirstChild();

		while( node ){
			int	parentId = node->getParentId();

			if( m_lookupTable.contains( parentId ) )
				m_lookupTable[parentId]->addChild( node );
			else
				delete node;

			node = m_orphans->takeFirstChild();
		}

		delete m_orphans;
		m_orphans = NULL;
	}

	int		nextTaxonId = 1;
	QString	sql = "SELECT MAX(taxon_id) FROM taxon;";

	if( !server.exec( sql ) )
		return false;

	QSqlQuery	query = server.getQuery();

	if( query.first() )
		nextTaxonId = query.value( 0 ).toInt() + 1;

	if( m_root )
		m_root->finalize( nextTaxonId );

	return true;
}

/**
**/
bool TaxonUpdateTree::loadFromDB( SqlServer &server )
{
	QString	sql( "SELECT taxon_id, ncbi_taxon_id, parent_taxon_id, node_rank, genetic_code, "
				"mito_genetic_code, left_value, right_value FROM taxon" );

	if( !server.exec( sql ) )
		return false;

	QSqlQuery	query = server.getQuery();

	if( !query.first() )
		return true;

	bool		status = true;
	Progress	progress( "Loading existing taxonomy definitions...", query.size() );

	QHash<int,TaxonUpdateNode*>	localLookup;

	do {
		int		localId = query.value( 0 ).toInt();
		int		taxId = query.value( 1 ).toInt();
		int		parentId = query.value( 2 ).toInt();
		QString	rank = query.value( 3 ).toString();
		int		genCodeId = query.value( 4 ).toInt();
		int		mitoGenCodeId = query.value( 5 ).toInt();
		int		left = query.value( 6 ).toInt();
		int		right = query.value( 7 ).toInt();

		TaxonUpdateNode	*node = new TaxonUpdateNode( localId, taxId, parentId, rank, genCodeId, mitoGenCodeId, left, right );

		if( localLookup.contains( parentId ) )
			localLookup[parentId]->addChild( node );
		else if( !m_root )
			m_root = node;
		else {
			if( !m_orphans )
				m_orphans = new TaxonUpdateNode;
			m_orphans->addChild( node );
		}

		localLookup[localId] = node;
		m_lookupTable[taxId] = node;

		progress.step();

		if( progress.isCanceled() ){
			delete m_root;
			m_root = NULL;
			delete m_orphans;
			m_orphans = NULL;
			status = false;
			break;
		}
	} while( query.next() );

	if( m_orphans ){
		TaxonUpdateNode	*node = m_orphans->takeFirstChild();

		while( node ){
			int	parentId = node->getLocalParentId();

			if( localLookup.contains( parentId ) )
				localLookup[parentId]->addChild( node );
			else
				delete node;

			node = m_orphans->takeFirstChild();
		}

		delete m_orphans;
		m_orphans = NULL;
	}

	progress.finish();

	return status;
}

/**
**/
bool TaxonUpdateTree::updateDB( SqlServer &server )
{
	if( !m_root )
		return true;

	Progress	progress( "Updating database entries...", m_root->getNumChildren( true ) + 1 );

	if( !m_root->updateDB( server, &progress ) )
		return false;

	progress.finish();

	Progress			progress2( "Updating taxonomy names...", m_names.count() );
	QString				sql( "INSERT INTO taxon_name (taxon_id, name, name_class) VALUES " );
	SqlInsertBuilder	insert( sql, &server );

	QHashIterator<int,TaxonUpdateName>	nameIt( m_names );

//	if( !server.exec( "ALTER TABLE taxon_name DISABLE KEYS" ) )
//		return false;

	while( nameIt.hasNext() ){
		nameIt.next();

		TaxonUpdateNode	*node = m_lookupTable.value( nameIt.key(), NULL );

		if( !node ){
			Core::warning( QString( "No taxon found for id: %1" ).arg( nameIt.key() ) );
			continue;
		}

		QListIterator< QPair<QString,QString> >	nIt( nameIt.value().getNameIterator() );

		while( nIt.hasNext() ){
			const QPair<QString,QString>	&name = nIt.next();

			insert.addValue( node->getLocalId() );
			insert.addValue( name.first );
			insert.addValue( name.second );

			if( !insert.addRow() )
				return false;
		}

		progress2.step();
		if( progress2.isCanceled() )
			return false;
	}

	if( !insert.flush() )
		return false;

//	if( !server.exec( "ALTER TABLE taxon_name ENABLE KEYS" ) )
//		return false;

	progress2.finish();

	return true;
}

/**
**/
void TaxonUpdateTree::print( QTextStream &os )
{
	if( m_root )
		m_root->print( os );
}
